#include "../src/em.hh"
#include "../src/distribution.cc"
#include "../src/nominal.hh"

#include <cassert>
#include <iostream>

namespace
{
using namespace playerpiano;

struct EmTestDatum 
{ 
  struct
    {
      float    alpha[4 * 5 * 5];
      float    logbeta;
      float    priorz[5];
      Rating   ratings[8];
      float    logbetamean;
      float    logbetastddev;
    }                   inputs;

  struct
    {
      float    q;
      float    logbeta;
      float    pz[5];
    }                   desired_output;
};


static void
test_em (void)
{
  EmTestDatum data[] = 
    {{{{-1., -0.931586, -0.568516, -0.773293, -0.11278, -0.888373, -1.,  -0.709909, -0.850158, -0.743908, -0.369839, -0.215757, -1.,  -0.125187, -0.816869, -0.507633, -0.667936, -0.660819, -1.,  -0.202639, -0.204111, -0.999327, -0.747347, -0.329849, -1., -1.,  -0.926238, -0.237544, -0.78334, -0.654589, -0.994651, -1., -0.669028,  -0.0100464, -0.541809, -0.106278, -0.959119, -1., -0.159888, -0.7979,  -0.73644, -0.743362, -0.0347012, -1., -0.981031, -0.228807,  -0.0754258, -0.373883, -0.778392, -1., -1., -0.0246958, -0.076099,  -0.626535, -0.448543, -0.0984582, -1., -0.838555, -0.843195,  -0.793954, -0.103807, -0.169527, -1., -0.833149, -0.252145,  -0.997529, -0.210408, -0.673261, -1., -0.454245, -0.261089,  -0.467045, -0.638559, -0.473214, -1., -1., -0.0322821, -0.39162,  -0.264677, -0.694821, -0.00758631, -1., -0.31552, -0.638142,  -0.246278, -0.909128, -0.476966, -1., -0.794946, -0.452324,  -0.805321, -0.307439, -0.961797, -1., -0.200179, -0.807793,  -0.0970316, -0.288536, -0.745934, -1.},  0., {0.453296, 0.370014, 0.350023, 0.72728,  0.485578}, {{0, 2}, {0, 1}, {1, 1}, {0, 3}, {2, 0}, {2, 4}, {3,  0}, {3, 4}}, 0.761633, 0.6147}, {-11.231628867152420859,  0.73797357790339544017, {0.827763, 0.00610845, 0.00498052,  0.097779,  0.0633688}}}, {{{-1., -0.577899, -0.506835, -0.922846, -0.747159,  -0.331621, -1., -0.597707, -0.44588, -0.952212, -0.879297, -0.792386,  -1., -0.138441, -0.990415, -0.679118, -0.984593, -0.0414098, -1.,  -0.701879, -0.933184, -0.43789, -0.411423, -0.0519018, -1., -1.,  -0.660464, -0.923468, -0.173057, -0.666602, -0.082565, -1.,  -0.416633, -0.25021, -0.919443, -0.750944, -0.818925, -1., -0.80433,  -0.967231, -0.871647, -0.0265394, -0.665889, -1., -0.976815,  -0.192529, -0.0419461, -0.624479, -0.274937, -1., -1., -0.259344,  -0.604057, -0.213055, -0.223035, -0.59888, -1., -0.680589,  -0.0399986, -0.556433, -0.516315, -0.263956, -1., -0.789788,  -0.63699, -0.765371, -0.445031, -0.985458, -1., -0.669759, -0.893724,  -0.418491, -0.31957, -0.692944, -1., -1., -0.701195, -0.376545,  -0.695091, -0.418007, -0.441851, -1., -0.772489, -0.482036,  -0.194972, -0.842971, -0.0919, -1., -0.442037, -0.638539, -0.326656,  -0.827944, -0.652249, -1., -0.00154913, -0.561285, -0.382913,  -0.666791, -0.33179, -1.},  0., {0.332439, 0.0355782, 0.652779, 0.361154,  0.0336343}, {{1, 2}, {1, 1}, {2, 4}, {1, 2}, {2, 0}, {2, 3}, {3,  1}, {1, 0}}, 0.412123, 0.34787}, {-11.741524673683423995,  0.34187609746382479776, {0.576981, 0.0090921, 0.263278, 0.13172,  0.0189298}}}, {{{-1., -0.220839, -0.524514, -0.815388, -0.170094,  -0.0258663, -1., -0.681543, -0.723488, -0.728057, -0.387327,  -0.354887, -1., -0.895544, -0.0758082, -0.385778, -0.793602,  -0.512631, -1., -0.409018, -0.0539882, -0.126041, -0.548209,  -0.0617968, -1., -1., -0.415142, -0.159675, -0.960332, -0.409667,  -0.194304, -1., -0.635161, -0.144945, -0.239573, -0.168437,  -0.953618, -1., -0.421457, -0.511516, -0.78111, -0.598731, -0.525913,  -1., -0.435708, -0.395332, -0.805129, -0.0132822, -0.02669, -1., -1.,  -0.341344, -0.679088, -0.465073, -0.964893, -0.926202, -1.,  -0.519413, -0.504741, -0.555226, -0.731898, -0.884252, -1.,  -0.359796, -0.315653, -0.563461, -0.930635, -0.938339, -1.,  -0.804138, -0.782351, -0.331904, -0.412427, -0.36843, -1., -1.,  -0.387018, -0.526775, -0.399144, -0.34174, -0.0456743, -1.,  -0.847688, -0.934071, -0.376847, -0.119473, -0.328275, -1., -0.42933,  -0.821621, -0.387574, -0.444022, -0.0695341, -1., -0.505967,  -0.824114, -0.513388, -0.131195, -0.70183, -1.},  0., {0.958237, 0.818516, 0.281232, 0.6666,  0.345255}, {{0, 4}, {2, 3}, {0, 4}, {3, 1}, {1, 2}, {1, 1}, {1,  3}, {2, 0}}, 0.345292, 0.680376}, {-12.149688662115401111,  0.19688070541945473773, {0.0623987, 0.656569, 0.0228284, 0.11295,  0.145253}}}, {{{-1., -0.99166, -0.60907, -0.807021, -0.385553,  -0.614813, -1., -0.489598, -0.478746, -0.956222, -0.793192,  -0.102023, -1., -0.0347236, -0.886688, -0.287225, -0.27791,  -0.521336, -1., -0.755494, -0.585396, -0.236147, -0.339852,  -0.0367256, -1., -1., -0.251996, -0.581402, -0.685144, -0.717102,  -0.260337, -1., -0.972332, -0.878123, -0.331549, -0.645524,  -0.482735, -1., -0.399377, -0.375327, -0.852332, -0.380711,  -0.364654, -1., -0.488638, -0.565107, -0.102801, -0.843318,  -0.733144, -1., -1., -0.979711, -0.866654, -0.503466, -0.696419,  -0.727715, -1., -0.285252, -0.818322, -0.979317, -0.467378, -0.31292,  -1., -0.940198, -0.647768, -0.821855, -0.830185, -0.540821, -1.,  -0.272442, -0.969523, -0.449474, -0.176167, -0.783803, -1., -1.,  -0.404416, -0.346673, -0.332849, -0.050659, -0.424705, -1.,  -0.480018, -0.829383, -0.35424, -0.69699, -0.194766, -1., -0.0110614,  -0.374923, -0.229612, -0.881846, -0.070863, -1., -0.727155,  -0.407757, -0.0516613, -0.530042, -0.454714, -1.},  0., {0.561766, 0.397813, 0.646125, 0.32909,  0.966182}, {{3, 2}, {1, 3}, {1, 2}, {3, 0}, {0, 2}, {1, 2}, {1,  0}, {1, 4}}, 0.744485, 0.978974}, {-11.425256244640923252,  0.81106631162037174171, {0.0922968, 0.00326753, 0.810237,  0.0011588, 0.0930394}}}};

  for (unsigned int i = 0; i < sizeof (data) / sizeof (data[0]); ++i)
    {
      NormalDistribution prior_logbeta (data[i].inputs.logbetamean, 
                                        data[i].inputs.logbetastddev);
      Nominal nominal (data[i].inputs.alpha, 0, 0, 4, 5);

      for (unsigned int j = 0; j < 4; ++j)
        {
          for (unsigned int k = 0; k < 5; ++k)
            {
              data[i].inputs.alpha[j * 5 * 5 + k * 5 + k] = drand48 ();
            }
        }

      std::vector<float> pz (5);
      std::vector<float> logpriorz (5);
      std::vector<Rating> ratings (8);

      for (unsigned int k = 0; k < 5; ++k)
        {
          logpriorz[k] = log (data[i].inputs.priorz[k]);
          pz[k] = drand48 ();
        }

      for (unsigned int k = 0; k < 8; ++k)
        {
          ratings[k] = data[i].inputs.ratings[k];
        }

      float logbeta = data[i].inputs.logbeta;

      float q = em (nominal,
                    logbeta,
                    pz,
                    logpriorz,
                    ratings,
                    prior_logbeta,
                    1000,
                    1e-7);

      assert (fabs (q - data[i].desired_output.q) <=
              1e-3 * (1 + fabs (q) + fabs (data[i].desired_output.q)) ||
              (std::cerr << q << " ?= " << data[i].desired_output.q << std::endl,
               0));

      assert (fabs (logbeta - data[i].desired_output.logbeta) <=
              5e-3 * (1 + fabs (logbeta) + fabs (data[i].desired_output.logbeta)) ||
              (std::cerr << logbeta << " ?= " << data[i].desired_output.logbeta << std::endl,
               0));

      for (unsigned int k = 0; k < 5; ++k)
        {
          assert (fabs (pz[k] - data[i].desired_output.pz[k]) <=
                  5e-3 * (1 + fabs (pz[k]) + fabs (data[i].desired_output.pz[k])) ||
                  (std::cerr << "pz[" << k << "] = " << pz[k] << " ?= " << data[i].desired_output.pz[k] << std::endl,
                   0));
        }
    }
}

}

int 
main (void) 
{
  srand48 (69);

  test_em ();

  return 0;
}
